#include "RegistryKey.h"

#include <iostream>

#ifndef KEY_WOW64_64KEY
// Access a 32-bit key from either a 32-bit or 64-bit application.
#define KEY_WOW64_64KEY 0x0100
#endif

#ifndef KEY_WOW64_32KEY
// Access a 64-bit key from either a 32-bit or 64-bit application.
#define KEY_WOW64_32KEY 0x0200
#endif

#define MAX_REGISTRY_LEN 4096

RegistryKey::RegistryKey(const std::string &registryPath, RegistryKey::Format fmt)
    : m_regKey(registryPath), m_fmt(fmt)
{
    m_errorCode = OpenKey(registryPath);
}

RegistryKey::RegistryKey(HKEY parentKey, const std::string &subKey, RegistryKey::Format fmt)
    : m_regKey(subKey), m_fmt(fmt)
{
    m_errorCode = OpenKey(parentKey, subKey);
}

RegistryKey::~RegistryKey()
{
    CloseKey();
}

bool RegistryKey::ReadSubKey(const std::string &subKey, unsigned long &value)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }
    DWORD dwType = REG_DWORD;
    DWORD dwSize = sizeof(DWORD);
    auto hr = RegQueryValueExA(m_hKey, subKey.c_str(), 0, &dwType, (LPBYTE)&value, &dwSize);
    SetError(hr);
    return hr == ERROR_SUCCESS;
}

bool RegistryKey::ReadSubKey(const std::string &subKey, std::string &value)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }
    char buf[MAX_REGISTRY_LEN] = {0};
    size_t dwSize = MAX_REGISTRY_LEN;
    DWORD dwType = REG_MULTI_SZ;
    auto hr = RegQueryValueExA(m_hKey, subKey.c_str(), 0, &dwType, (BYTE *)buf, (LPDWORD)&dwSize);
    value = std::string(buf);
    SetError(hr);
    return hr == ERROR_SUCCESS;
}

bool RegistryKey::ReadSubKey(const std::string &subKey, const char *buf, size_t &bufSize)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }
    DWORD dwType = REG_BINARY;
    auto hr = RegQueryValueExA(m_hKey, subKey.c_str(), 0, &dwType, (BYTE *)buf, (LPDWORD)&bufSize);
    SetError(hr);
    return hr == ERROR_SUCCESS;
}

// write key
bool RegistryKey::WriteSubKey(const std::string &subKey, const unsigned long &value)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }
    auto hr = RegSetValueExA(m_hKey, subKey.c_str(), 0, REG_DWORD, (LPCBYTE)&value, sizeof(value));
    SetError(hr);
    return hr == ERROR_SUCCESS;
}

bool RegistryKey::WriteSubKey(const std::string &subKey, const std::string &value)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }
    auto hr = RegSetValueExA(m_hKey, subKey.c_str(), 0, REG_MULTI_SZ, (LPCBYTE)value.c_str(), value.size());
    SetError(hr);
    return hr == ERROR_SUCCESS;
}

bool RegistryKey::WriteSubKey(const std::string &subKey, const char *buf, size_t bufSize)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }
    auto hr = RegSetValueExA(m_hKey, subKey.c_str(), 0, REG_BINARY, (LPCBYTE)buf, bufSize);
    SetError(hr);
    return hr == ERROR_SUCCESS;
}

bool RegistryKey::DeleteSubKey(const std::string &subKey)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }
    auto hr = RegDeleteValueA(m_hKey, subKey.c_str());
    SetError(hr);
    return hr == ERROR_SUCCESS;
}

bool RegistryKey::DeleteSubTree(const std::string &subKey)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }
    auto hr = RegDeleteKeyA(m_hKey, subKey.c_str());
    SetError(hr);
    return hr == ERROR_SUCCESS;
}

bool RegistryKey::AllChildKeys(std::vector<std::string> &results)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }

    DWORD numKeys;
    DWORD maxKeySize;
    DWORD numSubgroups;
    DWORD maxSubgroupSize;

    // Find the number of keys and subgroups, as well as the max of their lengths.
    LSTATUS hr = RegQueryInfoKey(m_hKey, 0, 0, 0, &numSubgroups, &maxSubgroupSize, 0, &numKeys,
                                 &maxKeySize, 0, 0, 0);
    if (hr != ERROR_SUCCESS)
    {
        SetError(hr);
        return false;
    }

    ++maxSubgroupSize;
    ++maxKeySize;

    DWORD n = numKeys;
    DWORD m = maxKeySize;
    ++m;

    char *buf = new char[m];
    memset(buf, '\0', m);
    for (int i = 0; i < n; ++i)
    {
        m = maxSubgroupSize;
        hr = RegEnumValueA(m_hKey, i, buf, &m, 0, 0, 0, 0);
        if (hr == ERROR_SUCCESS)
        {
            results.emplace_back(buf);
        }
        if (hr != ERROR_SUCCESS)
        {
            SetError(hr);
            continue;
        }
    }
    delete[] buf;

    if (results.empty() && m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }

    return true;
}

bool RegistryKey::AllChildGroups(std::vector<std::string> &results)
{
    if (m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }

    DWORD numKeys;
    DWORD maxKeySize;
    DWORD numSubgroups;
    DWORD maxSubgroupSize;

    // Find the number of keys and subgroups, as well as the max of their lengths.
    auto hr = RegQueryInfoKey(m_hKey, 0, 0, 0, &numSubgroups, &maxSubgroupSize, 0, &numKeys,
                              &maxKeySize, 0, 0, 0);
    if (hr != ERROR_SUCCESS)
    {
        SetError(hr);
        return false;
    }
    ++maxSubgroupSize;
    ++maxKeySize;

    DWORD n = numSubgroups;
    DWORD m = maxSubgroupSize;
    ++m;

    char *buf = new char[m];
    for (int i = 0; i < n; ++i)
    {
        m = maxSubgroupSize;
        memset(buf, '\0', m);
        hr = RegEnumKeyExA(m_hKey, i, buf, &m, 0, 0, 0, 0);
        if (hr == ERROR_SUCCESS)
        {
            results.emplace_back(buf);
        }

        if (hr != ERROR_SUCCESS)
        {
            SetError(hr);
            continue;
        }
    }
    delete[] buf;

    if (results.empty() && m_errorCode != E_REG_SUCCESS)
    {
        return false;
    }

    return true;
}

RegistryKey::ErrorCode RegistryKey::OpenKey(HKEY hParentKey, const std::string &subKey)
{
    LSTATUS hResult = ::RegOpenKeyExA(hParentKey, subKey.c_str(), 0, GetAcess(m_fmt), &m_hKey);
    SetError(hResult);
    if (ERROR_FILE_NOT_FOUND == hResult)
    {
        hResult = RegCreateKeyExA(hParentKey, m_regKey.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE,
                                  GetAcess(m_fmt), NULL, &m_hKey, NULL);

        SetError(hResult);
    }
    return m_errorCode;
}

RegistryKey::ErrorCode RegistryKey::OpenKey(const std::string &regPath)
{
    HKEY hParentKey{nullptr};
    if (GetHKeyInPath(m_regKey, hParentKey))
    {
        auto hResult = ::RegOpenKeyExA(hParentKey, m_regKey.c_str(), 0, GetAcess(m_fmt), &m_hKey);
        SetError(hResult);
        if (ERROR_FILE_NOT_FOUND == hResult)
        {
            if (hParentKey)
            {
                hResult = RegCreateKeyExA(hParentKey, m_regKey.c_str(), 0, NULL, REG_OPTION_NON_VOLATILE,
                                          GetAcess(m_fmt), NULL, &m_hKey, NULL);
                SetError(hResult);
            }
        }
    }
    return m_errorCode;
}

void RegistryKey::CloseKey()
{
    ::RegCloseKey(m_hKey);
}

bool RegistryKey::GetHKeyInPath(std::string &regPath, HKEY &key)
{
    // Function 'std::string_view::starts_with' requires C++ 20 standard
    int keyLength{0};

    std::string_view svRegPath(regPath);
    if (svRegPath.starts_with('\\'))
    {
        svRegPath.remove_prefix(1);
    }

    if (svRegPath.starts_with("HKEY_CURRENT_USER"))
    {
        keyLength = 17;
        key = HKEY_CURRENT_USER;
    }
    else if (svRegPath.starts_with("HKCU"))
    {
        keyLength = 4;
        key = HKEY_CURRENT_USER;
    }
    else if (svRegPath.starts_with("HKEY_LOCAL_MACHINE"))
    {
        keyLength = 18;
        key = HKEY_LOCAL_MACHINE;
    }
    else if (svRegPath.starts_with("HKLM"))
    {
        keyLength = 4;
        key = HKEY_LOCAL_MACHINE;
    }
    else if (svRegPath.starts_with("HKEY_CLASSES_ROOT"))
    {
        keyLength = 17;
        key = HKEY_CLASSES_ROOT;
    }
    else if (svRegPath.starts_with("HKCR"))
    {
        keyLength = 4;
        key = HKEY_CLASSES_ROOT;
    }
    else if (svRegPath.starts_with("HKEY_USERS"))
    {
        keyLength = 10;
        key = HKEY_USERS;
    }
    else if (svRegPath.starts_with("HKU"))
    {
        keyLength = 3;
        key = HKEY_USERS;
    }
    else
    {
        m_errorCode = E_REG_WRONG_PATH;
        return false;
    }

    svRegPath.remove_prefix(keyLength);
    if (svRegPath.starts_with('\\'))
    {
        svRegPath.remove_prefix(1);
    }

    regPath = svRegPath;
    return true;
}

unsigned long RegistryKey::GetAcess(RegistryKey::Format fmt)
{
    REGSAM access = KEY_READ | KEY_WRITE;
    switch (fmt)
    {
    case Registry32Format:
        access |= KEY_WOW64_32KEY;
        break;
    case Registry64Format:
        access |= KEY_WOW64_64KEY;
        break;
    default:
        break;
    }
    return access;
}

void RegistryKey::SetError(LSTATUS hResult)
{
    if (ERROR_SUCCESS == hResult)
    {
        m_errorCode = E_REG_SUCCESS;
    }
    else if (ERROR_ACCESS_DENIED == hResult)
    {
        m_errorCode = E_REG_NO_PERMISSION;
    }
    else if (ERROR_FILE_NOT_FOUND == hResult)
    {
        m_errorCode = E_REG_NOT_FIND;
    }
    else if (ERROR_INVALID_HANDLE == hResult)
    {
        m_errorCode = E_REG_WRONG_PATH;
    }
    else
    {
        m_errorCode = E_REG_UNKNOWN;
    }
}
